#include "memory.h"
#include "stdint.h"
#include "print.h"
#include "bitmap.h"
#include "global.h"
#include "debug.h"
#include "string.h"
#include "sync.h"

#define PG_SIZE 4096 //页大小 4K

#define MEM_BITMAP_BASE 0xc009a000

/*0xc000_0000 内核虚拟地址3G起
  0x0010_0000 跨过1M空间使虚拟地址在逻辑上连续*/
#define K_HEAP_START 0xc0100000

#define PDE_IDX(addr) ((addr & 0xffc00000) >> 22)
#define PTE_IDX(addr) ((addr & 0x003ff000) >> 12)

/*内存池结构，声明两个实例管理内核内存池和用户内存池*/
struct pool {

    struct lock lock; // 申请内存时互斥

    // 当前内存池的用到的位图结构，管理物理内存
    struct bitmap pool_bitmap;

    // 当前内存池管理的物理内存起始地址
    uint32_t phy_addr_start;

    //当前内存池的字节容量
    uint32_t pool_size; 
};
struct pool kernel_pool, user_pool;// 生成内核内存池，用户内存池
struct virtual_addr kernel_vaddr;// 内核分配虚拟地址


uint32_t* pte_ptr(uint32_t vaddr);
uint32_t* pde_ptr(uint32_t vaddr);
void* malloc_page(enum pool_flags pf, uint32_t pg_cnt);
void* get_user_pages(uint32_t pg_cnt);



/*在虚拟内存池中申请pg_cnt个内存，pf表示内核和用户
  成功返回虚拟页起始地址*/
static void* vaddr_get(enum pool_flags pf, uint32_t pg_cnt){
    int vaddr_start = 0, bit_idx_start = -1;
    uint32_t cnt = 0;
    if(pf == PF_KERNEL){
        //内核内存池
        bit_idx_start =bitmap_scan(&kernel_vaddr.vaddr_bitmap, pg_cnt);
        if(bit_idx_start == -1){
            return NULL;
        }
        while(cnt < pg_cnt){
            bitmap_set(&kernel_vaddr.vaddr_bitmap, bit_idx_start + cnt++, 1);
        }
        vaddr_start = kernel_vaddr.vaddr_start + bit_idx_start * PG_SIZE;
    } else {
        //用户内存池
        struct task_struct* cur = running_thread();
        bit_idx_start = bitmap_scan(&cur->userprog_vaddr.vaddr_bitmap, pg_cnt);
        if(bit_idx_start == -1){
            return NULL;
        }
        while(cnt < pg_cnt){
            bitmap_set(&cur->userprog_vaddr.vaddr_bitmap, bit_idx_start + cnt++, 1);
        }
        vaddr_start = cur->userprog_vaddr.vaddr_start + bit_idx_start * PG_SIZE;

        //(0xc0000000 - PG_SIZE)作为用户3级栈已经在start_process被分配
        ASSERT((uint32_t)vaddr_start < (0xc0000000 - PG_SIZE));
    }
    return (void*)vaddr_start;
}

/*得到虚拟地址vaddr对应的pte指针*/
uint32_t* pte_ptr(uint32_t vaddr){
    uint32_t* pte = (uint32_t*)(0xffc00000 + ((vaddr & 0xffc00000) >> 10) + PTE_IDX(vaddr) * 4);
    return pte;
}

/*得到虚拟地址vaddr对应的pde指针*/
uint32_t* pde_ptr(uint32_t vaddr){
    uint32_t* pde = (uint32_t*)((0xfffff000) + PDE_IDX(vaddr) * 4);
    return pde;
}

/*在m_pool指向的物理内存池中分配1个物理页
  成功则返回页框的物理地址*/
static void* palloc(struct pool* m_pool){
    int bit_idx = bitmap_scan(&m_pool->pool_bitmap, 1);
    if(bit_idx == -1){
        return NULL;
    }
    bitmap_set(&m_pool->pool_bitmap, bit_idx, 1);
    uint32_t page_phyaddr = ((bit_idx * PG_SIZE) + m_pool->phy_addr_start);
    return (void*)page_phyaddr;
}

//页表中添加虚拟地址_vaddr与物理地址_page_phyaddr的映射
static void page_table_add(void* _vaddr, void* _page_phyaddr){
    uint32_t vaddr = (uint32_t)_vaddr,page_phyaddr = (uint32_t)_page_phyaddr;
    uint32_t* pde = pde_ptr(vaddr);
    uint32_t* pte = pte_ptr(vaddr);
    //判断p位页目录项是否存在
    if(*pde & 0x00000001){
        ASSERT(!(*pte & 0x00000001));
        if(!(*pte & 0x00000001)){
            *pte = (page_phyaddr | PG_US_U | PG_RW_W | PG_P_1);
        }else{
            PANIC("pte repeat");
            *pte = (page_phyaddr | PG_US_U | PG_RW_W | PG_P_1);
        }
    }else{
        uint32_t pde_phyaddr = (uint32_t)palloc(&kernel_pool);
        *pde = (pde_phyaddr | PG_US_U | PG_RW_W | PG_P_1);
        memset((void*)((int)pte & 0xfffff000), 0, PG_SIZE);
        ASSERT(!(*pte & 0x00000001));
        *pte = (page_phyaddr | PG_US_U | PG_RW_W | PG_P_1);
    }
}

//分配pg_cnt个页空间，成功则返回虚拟地址
void* malloc_page(enum pool_flags pf, uint32_t pg_cnt){
    ASSERT(pg_cnt > 0 && pg_cnt < 3840);
    /*
        vaddr_get申请虚拟地址
        palloc申请物理页
        page_table_add将虚拟地址和物理地址在页表中映射
    */
    void* vaddr_start = vaddr_get(pf, pg_cnt);
    if(vaddr_start == NULL){
        return NULL;
    }
    uint32_t vaddr = (uint32_t)vaddr_start, cnt = pg_cnt;
    struct pool* mem_pool = pf & PF_KERNEL ? &kernel_pool : &user_pool;

    while(cnt-- > 0){
        void* page_phyaddr = palloc(mem_pool);
        if(page_phyaddr == NULL){
            return NULL;
        }
        page_table_add((void*)vaddr, page_phyaddr);
        vaddr += PG_SIZE;
    }
    return vaddr_start;
}

void* get_kernel_pages(uint32_t pg_cnt){
    void* vaddr = malloc_page(PF_KERNEL, pg_cnt);
    if(vaddr != NULL){
        memset(vaddr, 0, pg_cnt * PG_SIZE);
    }
    return vaddr;
}

/*在用户空间中申请4k内存，并返回虚拟地址*/
void* get_user_pages(uint32_t pg_cnt){
    lock_acquire(&user_pool.lock);
    void* vaddr = malloc_page(PF_USER, pg_cnt);
    memset(vaddr, 0, pg_cnt * PG_SIZE);
    lock_release(&user_pool.lock);
    return vaddr;
}

/*将地址vaddr与pf池中的物理地址关联，仅支持一页分配*/
void* get_a_page(enum pool_flags pf, uint32_t vaddr){
    struct pool* mem_pool = pf & PF_KERNEL ? &kernel_pool : &user_pool;
    lock_acquire(&mem_pool->lock);

    /*将虚拟地址对应的位图置1*/
    struct task_struct* cur = running_thread();
    int32_t bit_idx = -1;

    /*若当前是用户进程申请用户内存，就修改用户进程自己的虚拟地址位图*/
    if(cur->pgdir != NULL && pf == PF_USER){
        bit_idx = (vaddr - cur->userprog_vaddr.vaddr_start) / PG_SIZE;
        ASSERT(bit_idx > 0);
        bitmap_set(&cur->userprog_vaddr.vaddr_bitmap, bit_idx, 1);

    }else if(cur->pgdir == NULL && pf == PF_KERNEL){
        /*如果是内核线程申请内核内存，就修改kernel_vaddr*/
        bit_idx = (vaddr - kernel_vaddr.vaddr_start) / PG_SIZE;
        ASSERT(bit_idx > 0);
        bitmap_set(&kernel_vaddr.vaddr_bitmap, bit_idx, 1);

    }else{
        PANIC("get_a_page:not allow kernel alloc userspace or user alloc kernelspace by get_a_page");
    }
    void* page_phyaddr = palloc(mem_pool);
    if(page_phyaddr == NULL){
        return NULL;
    }
    page_table_add((void*)vaddr, page_phyaddr);
    lock_release(&mem_pool->lock);
    return (void*)vaddr;
}

/*将虚拟地址映射到物理地址*/
uint32_t addr_v2p(uint32_t vaddr){
    uint32_t* pte = pte_ptr(vaddr);
    /*(*pte)的值是页表所在的物理页框地址，
      去掉其低12位的页表项属性+虚拟地址vaddr的低12位*/
    return((*pte & 0xfffff000) + (vaddr & 0x00000fff));
}

// 内存池初始化
static void mem_pool_init(uint32_t all_mem){
    put_str("  mem_pool_init start\n");

    /*页表总大小 769-1022共254项，0和768项指向同一个页表
    ，一页的页目录，共256*页大小4k*/
    uint32_t page_table_size = PG_SIZE * 256;
    //已用内存数
    uint32_t used_mem = page_table_size + 0x100000;
    //可用内存数
    uint32_t free_mem = all_mem - used_mem;
    //可以内存页数
    uint16_t all_free_pages = free_mem / PG_SIZE;

    uint16_t kernel_free_pages = all_free_pages / 2;
    uint16_t user_free_pages = all_free_pages - kernel_free_pages;

    //简化位图操作不出来余数
    //kernel bitmap的长度，位图中的一位表示一页，以字节为单位
    uint32_t kbm_length = kernel_free_pages / 8;
    //user bitmap的长度
    uint32_t ubm_length = user_free_pages / 8;

    //kernel pool start 内核内存池的起始地址
    uint32_t kp_start = used_mem;
    //user pool start 用户内存池的起始地址
    uint32_t up_start = kp_start + kernel_free_pages * PG_SIZE;
    
    kernel_pool.phy_addr_start = kp_start;
    user_pool.phy_addr_start = up_start;
    kernel_pool.pool_size = kernel_free_pages * PG_SIZE;
    user_pool.pool_size = user_free_pages * PG_SIZE;
    kernel_pool.pool_bitmap.btmp_bytes_len = kbm_length;
    user_pool.pool_bitmap.btmp_bytes_len = ubm_length;

    //内核内存池在0xc009a000处
    kernel_pool.pool_bitmap.bits = (void*)MEM_BITMAP_BASE;
    //用户内存池跟在内核内存池后 
    user_pool.pool_bitmap.bits = (void*)(MEM_BITMAP_BASE + kbm_length);

    put_str("    kernel_pool_bitmap_start:");
    put_int((int)kernel_pool.pool_bitmap.bits);
    put_str(" kernel_pool.phy_addr_start:");
    put_int(kernel_pool.phy_addr_start);
    put_str("\n");
    put_str("    user_pool_bitmap_start:");
    put_int((int)user_pool.pool_bitmap.bits);
    put_str(" user_pool_phy_addr_start:");
    put_int(user_pool.phy_addr_start);
    put_str("\n");

    //将位图置0
    bitmap_init(&kernel_pool.pool_bitmap);
    bitmap_init(&user_pool.pool_bitmap);

    lock_init(&kernel_pool.lock);
    lock_init(&user_pool.lock);

    //维护内核堆的虚拟地址，所以要和内核内存池大小一致
    kernel_vaddr.vaddr_bitmap.btmp_bytes_len = kbm_length;

    //虚拟地址位图数组指向用户内存池之后的未使用的内存
    kernel_vaddr.vaddr_bitmap.bits = \
    (void*)(MEM_BITMAP_BASE + kbm_length + ubm_length);

    //虚拟内存池的起始地址0xc010_0000
    kernel_vaddr.vaddr_start = K_HEAP_START;
    //将位图置0
    bitmap_init(&kernel_vaddr.vaddr_bitmap);
    put_str("  mem_pool_init done\n");
}
void mem_init(void){
    put_str("mem_init start\n");

    //在loader中保存的内存大小在地址0xb00处
    uint32_t mem_bytes_total = *((uint32_t*)(0xb00));
    //put_int(mem_bytes_total);
    mem_pool_init(mem_bytes_total);
    put_str("mem_init done\n");
}

